from manim import *
import multiprocessing

COMPARE_POS1 = [-5, 3, 0]
COMPARE_POS = [-4, 3, 0]
COMPARE_POS2 = [-3, 3, 0]
GAP_POS = [-5, 3, 0]
SIDE_LENGTH = 1
POSX = -6
POSY = -3
POSZ = 0


class Element:
    value:int
    mob:any

    def __init__(self, value, mob):
        self.value = value
        self.mob = mob

class Sort(Scene):

    def construct(self):
        ValueArr = [5, 9, 1, 7, 12, 3, 10, 2, 8, 6, 4, 11]
        # ValueArr = [5, 9, 1, 7, 3, 10]
        # ValueArr = [5, 9, 1, 7]
        Arr = []
        i = 0
        while i < len(ValueArr):
            element_shape = Rectangle(width = SIDE_LENGTH,
                                      height = SIDE_LENGTH,
                                      color = "#808080",
                                      fill_color = WHITE,
                                      stroke_width=4,
                                      fill_opacity=1)
            element_text = Text(str(ValueArr[i]), font_size = 48, color = BLACK)
            element = VGroup(element_shape, element_text)
            element.move_to([POSX + i * SIDE_LENGTH, POSY, POSZ])
            Arr.append(Element(ValueArr[i], element))
            self.add(element)
            self.bring_to_front(element_text)
            self.wait(0.1)
            i += 1

        # self.BubbleSort(Arr)
        self.SelectSort(Arr)

    def Compare(self, e1, e2, swap):
        animations = []
        e1.mob[1].set_z_index(2)
        e2.mob[1].set_z_index(2)
        animations.append(e1.mob[0].animate.set_fill(RED,opacity=1).set_stroke(e1.mob[0].stroke_color))
        animations.append(e2.mob[0].animate.set_fill(RED,opacity=1).set_stroke(e2.mob[0].stroke_color))
        self.play(AnimationGroup(*animations, lag_ratio=0))
        self.wait()

        string = str(e1.value)
        if e1.value > e2.value:
            string += " > "
        elif e1.value < e2.value:
            string += " < "
        else:
            string += " = "
        string += str(e2.value)

        if swap:
            string += ", 交换"
            # self.Swap(e1 ,e2)
        else:
            string += ", 不交换"

        text = Text(string, color = RED)
        text.move_to(COMPARE_POS)
        self.play(Create(text))
        self.wait(0.5)
        self.play(FadeOut(text))
        animations = []
        e1.mob[1].set_z_index(2)
        e2.mob[1].set_z_index(2)
        animations.append(e1.mob[0].animate.set_fill(WHITE,opacity=1).set_stroke(e1.mob[0].stroke_color))
        animations.append(e2.mob[0].animate.set_fill(WHITE,opacity=1).set_stroke(e2.mob[0].stroke_color))
        self.play(AnimationGroup(*animations, lag_ratio=0))

    def Swap(self, e1, e2):
        POS1 = e1.mob.get_center()
        POS2 = e2.mob.get_center()
        self.play(e1.mob.animate.move_to(POS2),
                  e2.mob.animate.move_to(POS1))
        
    def BubbleSort(self, Arr):
        n = len(Arr)
        i = 0
        j = 0
        textj = Text("j", color = YELLOW)
        textj.move_to([POSX, POSY + 1, POSZ])
        textj2 = Text("j+1", color = YELLOW)
        textj2.move_to([POSX + SIDE_LENGTH, POSY + 1, POSZ])
        self.add(textj)
        self.add(textj2)

        for i in range(n - 1):        
            for j in range(0, n - i - 1):
                self.play(textj.animate.move_to([POSX + (j) * SIDE_LENGTH, POSY + 1, POSZ]),
                          textj2.animate.move_to([POSX + (j + 1) * SIDE_LENGTH, POSY + 1, POSZ]))
                if Arr[j].value > Arr[j + 1].value:
                    self.Compare(Arr[j], Arr[j + 1], 1)
                    self.Swap(Arr[j], Arr[j + 1])
                    temp = Arr[j + 1]
                    Arr[j + 1] = Arr[j]
                    Arr[j] = temp
                else:
                    self.Compare(Arr[j], Arr[j + 1], 0)
                    animations = []
                self.play(textj.animate.move_to([POSX + (j) * SIDE_LENGTH, POSY + 1, POSZ]),
                          textj2.animate.move_to([POSX + (j + 1) * SIDE_LENGTH, POSY + 1, POSZ]))
            e1 = Arr[n - i - 1]
            e1.mob[1].set_z_index(2)
            animations.append(e1.mob[0].animate.set_fill(GREEN,opacity=1).set_stroke(e1.mob[0].stroke_color))
            self.play(AnimationGroup(*animations, lag_ratio=0))
        self.wait()
        e1 = Arr[0]
        e1.mob[1].set_z_index(2)
        animations.append(e1.mob[0].animate.set_fill(GREEN,opacity=1).set_stroke(e1.mob[0].stroke_color))
        self.play(AnimationGroup(*animations, lag_ratio=0))
        self.wait()
        self.play(FadeOut(textj), FadeOut(textj2))
        self.wait()

    def Compare2(self, e1, e2, index, stri):
        animations = []
        e1.mob[1].set_z_index(2)
        e2.mob[1].set_z_index(2)
        animations.append(e1.mob[0].animate.set_fill(RED,opacity=1).set_stroke(e1.mob[0].stroke_color))
        animations.append(e2.mob[0].animate.set_fill(RED,opacity=1).set_stroke(e2.mob[0].stroke_color))
        self.play(AnimationGroup(*animations, lag_ratio=0))
        self.wait()

        string = str(e1.value)
        if e1.value > e2.value:
            string += " > "
        elif e1.value < e2.value:
            string += " < "
        else:
            string += " = "
        string += str(e2.value)

        if index:
            string += ", 移动 "
            string += stri
            string += " 到 j"
        else:
            string += ", 不移动"
            string += stri
        text = Text(string, color = RED)
        text.move_to([-2, 3, 0])
        self.play(Create(text))
        self.wait(0.5)
        self.play(FadeOut(text))
        animations = []
        e1.mob[1].set_z_index(2)
        e2.mob[1].set_z_index(2)
        animations.append(e1.mob[0].animate.set_fill(WHITE,opacity=1).set_stroke(e1.mob[0].stroke_color))
        animations.append(e2.mob[0].animate.set_fill(WHITE,opacity=1).set_stroke(e2.mob[0].stroke_color))
        self.play(AnimationGroup(*animations, lag_ratio=0))


    def SelectSort(self, Arr):
        n = len(Arr)
        i = 0
        j = 0
        min_i = 0
        max_i = 0
        texti = Text("i", color = YELLOW)
        textj = Text("j", color = YELLOW)
        textmini = Text("min_i", color = YELLOW)
        textmaxi = Text("max_i", color = YELLOW)
        texti.move_to([POSX, POSY + 1, POSZ])
        textj.move_to([POSX, POSY + 2, POSZ])
        textmini.move_to([POSX, POSY + 3, POSZ])
        textmaxi.move_to([POSX, POSY + 4, POSZ])
        self.add(texti)
        self.add(textj)
        self.add(textmini)
        self.add(textmaxi)
        self.play(Create(texti), Create(textj), Create(textmini), Create(textmaxi))
        while i < n / 2:
            min_i = max_i = i
            j = i + 1
            if i > 0:
                tempstr = "第 " + str(i) + " 趟循环完成\n"
                tempstr += "重置 mini 和 maxi 为 i"
                texttemp = Text(tempstr)
                texttemp.move_to([-2, 3, 0])
                self.play(Create(texttemp))
                self.play(FadeOut(texttemp))
            self.play(texti.animate.move_to([POSX + (i) * SIDE_LENGTH, POSY + 1, POSZ]),
                      textmini.animate.move_to([POSX + (min_i) * SIDE_LENGTH, POSY + 3, POSZ]),
                      textmaxi.animate.move_to([POSX + (max_i) * SIDE_LENGTH, POSY + 4, POSZ]))     
            while j < n - i:
                self.play(textj.animate.move_to([POSX + (j) * SIDE_LENGTH, POSY + 2, POSZ]))
                if Arr[j].value < Arr[min_i].value:
                    self.Compare2(Arr[j], Arr[min_i], 1, "min_i")
                    min_i = j
                    self.play(textmini.animate.move_to([POSX + (min_i) * SIDE_LENGTH, POSY + 3, POSZ]))
                else:
                    self.Compare2(Arr[j], Arr[min_i], 0, "min_i")
                if Arr[j].value > Arr[max_i].value:
                    self.Compare2(Arr[j], Arr[max_i], 1, "max_i")
                    max_i = j
                    self.play(textmaxi.animate.move_to([POSX + (max_i) * SIDE_LENGTH, POSY + 4, POSZ]))
                else:
                    self.Compare2(Arr[j], Arr[max_i], 0, "max_i")
                j += 1
            if max_i == i:
                texttemp = Text("由于 max_i == i\n重新赋值 max_i = min_i")
                texttemp.move_to([-2, 3, 0])
                self.play(Create(texttemp))
                self.play(FadeOut(texttemp))
                max_i = min_i
                self.play(textmaxi.animate.move_to([POSX + (max_i) * SIDE_LENGTH, POSY + 4, POSZ]))
            texttemp = Text("交换 Arr[i] 和 Arr[min_i]")
            texttemp.move_to([-2, 3, 0])
            self.play(Create(texttemp))
            self.play(FadeOut(texttemp))
            self.Swap(Arr[i], Arr[min_i])
            self.play(Arr[min_i].mob[0].animate.set_fill(GREEN,opacity=1).set_stroke(Arr[min_i].mob[0].stroke_color))
            temp = Arr[i]
            Arr[i] = Arr[min_i]
            Arr[min_i] = temp


            
            texttemp = Text("安放Arr[max_i]")
            texttemp.move_to([-2, 3, 0])
            self.play(Create(texttemp))
            self.play(FadeOut(texttemp))
            self.Swap(Arr[n - i - 1], Arr[max_i])
            self.play(Arr[max_i].mob[0].animate.set_fill(GREEN,opacity=1).set_stroke(Arr[max_i].mob[0].stroke_color))
            temp = Arr[n - i - 1]
            Arr[n - i - 1] = Arr[max_i]
            Arr[max_i] = temp
            i += 1
        self.play(FadeOut(texti), FadeOut(textj), FadeOut(textmini), FadeOut(textmaxi))
        self.wait()